from openai import OpenAI

from prompts import system_prompt
from actions import available_actions
from helpers import extract_json

class Agent:
    def __init__(self, client: OpenAI,model:str):
        self.client = client
        self.model = model
        self.messages: list = []
        self.messages.append({"role": "system", "content": system_prompt})
        self.max_turns = 5

    def execute(self):
        completion = self.client.chat.completions.create(
            model=self.model,
            messages=self.messages
        )
        return completion.choices[0].message.content

    def run(self, message: str = ""):
        turn_count = 1
        self.messages.append({"role": "user", "content": message})

        while turn_count < self.max_turns:
            print(f"------Loop: {turn_count}")
            turn_count += 1

            response = self.execute()
            print(response)
            self.messages.append({"role": "assistant", "content": response})
            
            if '行动' in response:
                # 解析出Action
                json_function = extract_json(response)
                if not json_function:
                    raise Exception(f"Invalid JSON in response: {response}")
                
                function_name = json_function[0]['function_name']
                function_parms = json_function[0]['function_parms']
                if function_name not in available_actions:
                    raise Exception(f"Unknown action: {function_name}: {function_parms}")
                                
                # 调用Action
                print(f"***Running action: {function_name} {function_parms}")
                action_function = available_actions[function_name]
                result = action_function(**function_parms)
                function_result_message = f"行动响应: {result}"
                self.messages.append({"role": "user", "content": function_result_message})
                print(f"***Running result: {function_result_message}")
                continue

            if '答案' in response:
                break